#!/usr/bin/env python3
"""
Remote Control System - Capture Side
Monitors keyboard and mouse events and sends them via UDP
"""

import json
import socket
import time
import threading
import sys
from datetime import datetime
from pynput import mouse, keyboard
from colorama import init, Fore, Style

# Initialize colorama for colored output
init()

class InputCapture:
    def __init__(self, target_host='localhost', target_port=3001):
        self.target_host = target_host
        self.target_port = target_port
        self.is_capturing = False
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.event_count = 0
        
        # Event listeners
        self.mouse_listener = None
        self.keyboard_listener = None
        
        print(f"{Fore.CYAN}🎯 Input Capture initialized - Target: {target_host}:{target_port}{Style.RESET_ALL}")
    
    def start(self):
        """Start capturing input events"""
        if self.is_capturing:
            print(f"{Fore.YELLOW}⚠️  Already capturing input{Style.RESET_ALL}")
            return
        
        self.is_capturing = True
        print(f"{Fore.GREEN}🚀 Starting input capture...{Style.RESET_ALL}")
        
        # Start mouse listener
        self.mouse_listener = mouse.Listener(
            on_move=self.on_mouse_move,
            on_click=self.on_mouse_click,
            on_scroll=self.on_mouse_scroll
        )
        
        # Start keyboard listener
        self.keyboard_listener = keyboard.Listener(
            on_press=self.on_key_press,
            on_release=self.on_key_release
        )
        
        # Start listeners
        self.mouse_listener.start()
        self.keyboard_listener.start()
        
        print(f"{Fore.GREEN}✅ Input capture started successfully{Style.RESET_ALL}")
        print(f"{Fore.BLUE}📝 Monitoring keyboard and mouse events...{Style.RESET_ALL}")
        print(f"{Fore.BLUE}⏹️  Press Ctrl+C to stop{Style.RESET_ALL}")
        
        # Start statistics thread
        self.stats_thread = threading.Thread(target=self.show_stats, daemon=True)
        self.stats_thread.start()
    
    def stop(self):
        """Stop capturing input events"""
        if not self.is_capturing:
            print(f"{Fore.YELLOW}⚠️  Not currently capturing{Style.RESET_ALL}")
            return
        
        self.is_capturing = False
        
        if self.mouse_listener:
            self.mouse_listener.stop()
        if self.keyboard_listener:
            self.keyboard_listener.stop()
        
        self.socket.close()
        print(f"{Fore.RED}🛑 Input capture stopped{Style.RESET_ALL}")
        print(f"{Fore.CYAN}📈 Total events captured: {self.event_count}{Style.RESET_ALL}")
    
    def send_event(self, event_data):
        """Send event data via UDP"""
        if not self.is_capturing:
            return

        try:
            event_data['timestamp'] = time.time()
            message = json.dumps(event_data)
            self.socket.sendto(message.encode('utf-8'), (self.target_host, self.target_port))
            self.event_count += 1
            
            # Optional: Log events (comment out for better performance)
            # print(f"📤 Sent: {event_data['type']}")
            
        except Exception as e:
            print(f"{Fore.RED}❌ Error sending event: {e}{Style.RESET_ALL}")
    
    def on_mouse_move(self, x, y):
        """Handle mouse move events"""
        self.send_event({
            'type': 'mousemove',
            'x': x,
            'y': y
        })
    
    def on_mouse_click(self, x, y, button, pressed):
        """Handle mouse click events"""
        event_type = 'mousedown' if pressed else 'mouseup'
        button_name = button.name  # 'left', 'right', 'middle'
        
        self.send_event({
            'type': event_type,
            'x': x,
            'y': y,
            'button': button_name
        })
    
    def on_mouse_scroll(self, x, y, dx, dy):
        """Handle mouse scroll events"""
        self.send_event({
            'type': 'mousewheel',
            'x': x,
            'y': y,
            'dx': dx,
            'dy': dy
        })
    
    def on_key_press(self, key):
        """Handle key press events"""
        try:
            key_name = key.char if hasattr(key, 'char') and key.char else str(key)
            self.send_event({
                'type': 'keydown',
                'key': key_name
            })
        except Exception as e:
            print(f"{Fore.RED}❌ Error handling key press: {e}{Style.RESET_ALL}")
    
    def on_key_release(self, key):
        """Handle key release events"""
        try:
            key_name = key.char if hasattr(key, 'char') and key.char else str(key)
            self.send_event({
                'type': 'keyup',
                'key': key_name
            })
            
            # Stop on Ctrl+C
            if key == keyboard.Key.esc:
                print(f"\n{Fore.YELLOW}🔄 Escape key pressed, stopping capture...{Style.RESET_ALL}")
                self.stop()
                return False
                
        except Exception as e:
            print(f"{Fore.RED}❌ Error handling key release: {e}{Style.RESET_ALL}")
    
    def show_stats(self):
        """Show statistics every 10 seconds"""
        while self.is_capturing:
            time.sleep(10)
            if self.is_capturing:
                print(f"{Fore.CYAN}📊 Events captured: {self.event_count}{Style.RESET_ALL}")

def main():
    """Main function for CLI usage"""
    import argparse
    
    parser = argparse.ArgumentParser(description='Remote Control - Input Capture')
    parser.add_argument('host', nargs='?', default='localhost', help='Target host (default: localhost)')
    parser.add_argument('port', nargs='?', type=int, default=3001, help='Target port (default: 3001)')
    
    args = parser.parse_args()
    
    print(f"{Fore.MAGENTA}🖱️  Remote Control - Input Capture{Style.RESET_ALL}")
    print(f"{Fore.MAGENTA}====================================={Style.RESET_ALL}")
    print()
    
    capture = InputCapture(args.host, args.port)
    
    try:
        capture.start()
        
        # Keep the main thread alive
        while capture.is_capturing:
            time.sleep(1)
            
    except KeyboardInterrupt:
        print(f"\n{Fore.YELLOW}🔄 Shutting down capture...{Style.RESET_ALL}")
        capture.stop()
    except Exception as e:
        print(f"{Fore.RED}💥 Error: {e}{Style.RESET_ALL}")
        capture.stop()

if __name__ == '__main__':
    main()